#ifndef _TASKSYS_H
#define _TASKSYS_H

#include "itasksys.h"
#include "../common/CycleTimer.h"
#include <iostream>
#include <vector>
#include <queue>
#include <mutex>
#include <memory>
#include <condition_variable>
#include <string>
#include <thread>
#include <functional>
#include <type_traits>
#include <future>
#include <deque>
#include <unordered_map>

/*
 * TaskSystemSerial: This class is the student's implementation of a
 * serial task execution engine.  See definition of ITaskSystem in
 * itasksys.h for documentation of the ITaskSystem interface.
 */
class TaskSystemSerial: public ITaskSystem {
    public:
        TaskSystemSerial(int num_threads);
        ~TaskSystemSerial();
        const char* name();
        void run(IRunnable* runnable, int num_total_tasks);
        TaskID runAsyncWithDeps(IRunnable* runnable, int num_total_tasks,
                                const std::vector<TaskID>& deps);
        void sync();
};

/*
 * TaskSystemParallelSpawn: This class is the student's implementation of a
 * parallel task execution engine that spawns threads in every run()
 * call.  See definition of ITaskSystem in itasksys.h for documentation
 * of the ITaskSystem interface.
 */
class TaskSystemParallelSpawn: public ITaskSystem {
    public:
        TaskSystemParallelSpawn(int num_threads);
        ~TaskSystemParallelSpawn();
        const char* name();
        void run(IRunnable* runnable, int num_total_tasks);
        TaskID runAsyncWithDeps(IRunnable* runnable, int num_total_tasks,
                                const std::vector<TaskID>& deps);
        void sync();

        int num_threads;
};

/*
 * TaskSystemParallelThreadPoolSpinning: This class is the student's
 * implementation of a parallel task execution engine that uses a
 * thread pool. See definition of ITaskSystem in itasksys.h for
 * documentation of the ITaskSystem interface.
 */
class TaskSystemParallelThreadPoolSpinning: public ITaskSystem {
    public:
        TaskSystemParallelThreadPoolSpinning(int num_threads);
        ~TaskSystemParallelThreadPoolSpinning();
        const char* name();
        void run(IRunnable* runnable, int num_total_tasks);
        TaskID runAsyncWithDeps(IRunnable* runnable, int num_total_tasks,
                                const std::vector<TaskID>& deps);
        void sync();

        public:
        // Enqueue A Task
        // template<typename F, typename ... Args>
        // auto enqueue(F&& f, Args&& ... args) -> std::future<typename std::result_of<F(Args...)>::type>;

        template<typename F, typename ... Args>
        auto enqueue(F&& f, Args&& ... args)
        -> std::future<typename std::result_of<typename std::decay<F>::type(typename std::decay<Args>::type ...)>::type>;

    private:
        // number of threads
        int num_threads;
        // thread pool
        std::vector<std::thread> workers;
        // task queue
        std::queue<std::function<void()>> task_queue;
        // synchronization - spinning lock
        std::mutex queue_mutex;
        // whether to stop the thread pool
        bool stop;
};

struct TaskInfo
{
    TaskID id;
    IRunnable* runnable;
    int num_total_tasks;
    std::vector<TaskID> dependency;

    TaskInfo(TaskID id, IRunnable* irunnable, int num_total_tasks, 
    const std::vector<TaskID>& depends) : id(id), runnable(irunnable), num_total_tasks(num_total_tasks),
    dependency(depends) {}
};

struct CompareTask {
    bool operator()(const std::pair<TaskID, TaskInfo>& a, const std::pair<TaskID, TaskInfo>& b) {
         if (a.second.id == b.second.id) {
            return a.second.num_total_tasks < b.second.num_total_tasks; 
        }
        return a.second.id > b.second.id; // 从小到大
    }
};

/*
 * TaskSystemParallelThreadPoolSleeping: This class is the student's
 * optimized implementation of a parallel task execution engine that uses
 * a thread pool. See definition of ITaskSystem in
 * itasksys.h for documentation of the ITaskSystem interface.
 */
class TaskSystemParallelThreadPoolSleeping: public ITaskSystem {
    public:
        TaskSystemParallelThreadPoolSleeping(int num_threads);
        ~TaskSystemParallelThreadPoolSleeping();
        const char* name();
        void run(IRunnable* runnable, int num_total_tasks);
        TaskID runAsyncWithDeps(IRunnable* runnable, int num_total_tasks,
                                const std::vector<TaskID>& deps);
        void sync();

    public:
        // std::mutex completed_mutex;
        // std::unordered_map<TaskID, std::atomic_bool> completed;

        struct subTaskManage
        {
            std::atomic_int64_t num_sub_tasks;
            std::atomic_bool is_done;

            subTaskManage() : num_sub_tasks(0), is_done(false) {}

            subTaskManage(int64_t num_sub_tasks, bool is_done) : num_sub_tasks(num_sub_tasks), is_done(is_done) {}

            subTaskManage(const subTaskManage&) = delete; 
            subTaskManage& operator=(const subTaskManage&) = delete; 

            subTaskManage(subTaskManage&& subTaskManage) : num_sub_tasks(std::move(subTaskManage.num_sub_tasks.load())), is_done(std::move(subTaskManage.is_done.load())) {}
            subTaskManage& operator=(subTaskManage&& subTaskManage) {
                if (this != &subTaskManage) {
                    num_sub_tasks = std::move(subTaskManage.num_sub_tasks.load());
                    is_done = std::move(subTaskManage.is_done.load());
                }
                return *this;
            }
        };

        std::mutex subTask_mutex;
        std::unordered_map<TaskID, subTaskManage> subTask;

    private:
        // Task Grunuality
        static const int task_granularity = 8; // each thread to process task nums is 32

        int num_threads;
        bool stop;
        std::vector<std::thread> workers;
        std::queue<std::function<void()>> task_queue;
        std::mutex queue_mutex;
        std::condition_variable condition;

        //  Sync
        std::mutex waiting_queue_mutex;
        //  std::mutex completed_mutex;
        std::mutex future_mutex;
        std::atomic<TaskID> id;

        // std::deque<std::pair<TaskID, TaskInfo>> waiting_queue;
        // Num_total_task priority_queue
        std::priority_queue<std::pair<TaskID, TaskInfo>, std::vector<std::pair<TaskID, TaskInfo>>, CompareTask> waiting_queue;
        std::condition_variable waiting_queue_condition;
        // std::vector<std::future<void>> futures;
        std::unordered_map<TaskID, std::vector<std::future<void>>> subTask_futures;

        template<typename F, typename ... Args>
        auto enqueue(F&& f, Args&& ... args)
        -> std::future<typename std::result_of<typename std::decay<F>::type(typename std::decay<Args>::type...)>::type>;

        void check_waiting_queue();

};

 
#endif
